package com.demo.wormhole.component.rabbit.factory;

import com.demo.wormhole.core.context.WormholeContextHolder;
import org.apache.commons.lang3.StringUtils;
import org.springframework.amqp.rabbit.config.SimpleRabbitListenerContainerFactory;
import org.springframework.amqp.rabbit.listener.SimpleMessageListenerContainer;

public class WormholeSimpleRabbitListenerContainerFactory extends SimpleRabbitListenerContainerFactory {
    
    /**
     * 接收消息的时候,需要根据消息的token设置ThreadLocal
     *
     * @return
     */
    @Override
    protected SimpleMessageListenerContainer createContainerInstance() {
        SimpleMessageListenerContainer simpleMessageListenerContainer = new SimpleMessageListenerContainer();
        
        simpleMessageListenerContainer.setAfterReceivePostProcessors(message -> {
            //防止线程复用 销毁ThreadLocal
            WormholeContextHolder.invalidContext();
            //获取消息属性标识
            String wormholeRequestContext = message
                    .getMessageProperties()
                    .getHeader(WormholeContextHolder.WORMHOLE_REQUEST_MARK);
            
            //取出消息的时候,需要做的操作,保证染色成功
            if (StringUtils.isNotEmpty(wormholeRequestContext)) {
                WormholeContextHolder.setContext(wormholeRequestContext);
            }
            return message;
        });
        
        return simpleMessageListenerContainer;
    }
}
